import cv2
import numpy as np
import tkinter as tk
from tkinter import filedialog, Scale, Button, Label, Frame, IntVar, StringVar, DoubleVar, Canvas
from PIL import Image, ImageTk
import os
import time

class ImageProcessor:
    """核心图像处理类，提供各种图像处理方法"""
    
    def __init__(self):
        self.image = None
        self.original = None
        self.mask = None
        self.height = 0
        self.width = 0
        self.selected_color = None  # 存储选中的颜色
        
    def load_image(self, image_path):
        """加载图像并保存原始副本"""
        self.image = cv2.imread(image_path)
        if self.image is None:
            raise ValueError(f"无法读取图像: {image_path}")
        self.original = self.image.copy()
        self.height, self.width = self.image.shape[:2]
        return self.image
        
    def preprocess_image(self, blur_size=5, brightness=0, contrast=1.0):
        """图像预处理：灰度转换、亮度对比度调整、高斯模糊"""
        if self.image is None:
            raise ValueError("请先加载图像")
            
        # 转换为灰度图像
        gray = cv2.cvtColor(self.image, cv2.COLOR_BGR2GRAY)
        
        # 调整亮度和对比度
        adjusted = cv2.addWeighted(gray, contrast, np.zeros_like(gray), 0, brightness)
        
        # 应用高斯模糊
        if blur_size > 0:
            blur_size = blur_size if blur_size % 2 == 1 else blur_size + 1  # 确保是奇数
            blurred = cv2.GaussianBlur(adjusted, (blur_size, blur_size), 0)
        else:
            blurred = adjusted
            
        return blurred
        
    def enhance_contrast(self, img, clip_limit=2.0, tile_grid_size=(8, 8)):
        """使用CLAHE增强对比度"""
        clahe = cv2.createCLAHE(clipLimit=clip_limit, tileGridSize=tile_grid_size)
        enhanced = clahe.apply(img)
        return enhanced
        
    def adaptive_threshold(self, img, block_size=11, c=2):
        """自适应阈值处理"""
        # 确保block_size是奇数
        block_size = block_size if block_size % 2 == 1 else block_size + 1
        
        # 应用自适应阈值
        thresh = cv2.adaptiveThreshold(
            img, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, 
            cv2.THRESH_BINARY_INV, block_size, c
        )
        return thresh
        
    def otsu_threshold(self, img):
        """使用Otsu算法进行阈值处理"""
        _, thresh = cv2.threshold(img, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
        return thresh
        
    def apply_morphology(self, img, operation='open', kernel_size=3):
        """应用形态学操作"""
        kernel = np.ones((kernel_size, kernel_size), np.uint8)
        
        if operation == 'dilate':
            result = cv2.dilate(img, kernel, iterations=1)
        elif operation == 'erode':
            result = cv2.erode(img, kernel, iterations=1)
        elif operation == 'open':
            result = cv2.morphologyEx(img, cv2.MORPH_OPEN, kernel)
        elif operation == 'close':
            result = cv2.morphologyEx(img, cv2.MORPH_CLOSE, kernel)
        else:
            result = img
            
        return result
        
    def filter_contours(self, img, min_area=100):
        """基于面积过滤连通区域"""
        # 寻找所有轮廓
        contours, _ = cv2.findContours(img, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        
        # 创建空白蒙版
        mask = np.zeros_like(img)
        
        # 将大于最小面积的轮廓填充到蒙版上
        for contour in contours:
            if cv2.contourArea(contour) > min_area:
                cv2.drawContours(mask, [contour], -1, 255, -1)
                
        return mask
        
    def create_mask(self, blur_size=5, brightness=0, contrast=1.0, 
                    use_clahe=True, clip_limit=2.0, 
                    thresh_method='adaptive', block_size=11, c=2,
                    morph_op='close', kernel_size=3, min_area=100):
        """创建蒙版的完整流程"""
        if self.original is None:
            raise ValueError("请先加载图像")
        
        # 重置图像为原始状态
        self.image = self.original.copy()
        
        # 预处理
        processed = self.preprocess_image(blur_size, brightness, contrast)
        
        # 增强对比度
        if use_clahe:
            processed = self.enhance_contrast(processed, clip_limit)
        
        # 阈值处理
        if thresh_method == 'adaptive':
            thresh = self.adaptive_threshold(processed, block_size, c)
        else:
            thresh = self.otsu_threshold(processed)
        
        # 形态学操作
        morph = self.apply_morphology(thresh, morph_op, kernel_size)
        
        # 区域过滤
        if min_area > 0:
            mask = self.filter_contours(morph, min_area)
        else:
            mask = morph
        
        self.mask = mask
        return mask
        
    def create_color_mask(self, selected_color, tolerance=30):
        """基于选择的颜色创建蒙版"""
        if self.original is None:
            raise ValueError("请先加载图像")
            
        # 重置图像为原始状态
        self.image = self.original.copy()
        
        # 创建蒙版
        mask = np.zeros((self.height, self.width), dtype=np.uint8)
        
        # 计算颜色差异
        b, g, r = cv2.split(self.image)
        b_diff = np.abs(b.astype(np.int16) - selected_color[0])
        g_diff = np.abs(g.astype(np.int16) - selected_color[1])
        r_diff = np.abs(r.astype(np.int16) - selected_color[2])
        
        # 计算总差异
        color_diff = b_diff + g_diff + r_diff
        
        # 应用容差阈值生成蒙版
        mask[color_diff <= tolerance*3] = 255
        
        self.mask = mask
        return mask
        
    def save_mask(self, output_path):
        """保存蒙版图像"""
        if self.mask is None:
            raise ValueError("请先创建蒙版")
        
        cv2.imwrite(output_path, self.mask)
        return output_path


class MaskGeneratorUI:
    """用户界面类，提供交互调整参数的界面"""
    
    def __init__(self, root):
        self.root = root
        self.root.title("艺术字蒙版生成器")
        self.root.geometry("1000x700")
        
        self.processor = ImageProcessor()
        self.image_path = ""
        self.color_pick_mode = False  # 颜色拾取模式标志
        self.setup_ui()
        
    def setup_ui(self):
        """设置UI组件"""
        # 创建左右布局框架
        left_frame = Frame(self.root)
        left_frame.pack(side=tk.LEFT, fill=tk.BOTH, expand=True)
        
        right_frame = Frame(self.root)
        right_frame.pack(side=tk.RIGHT, fill=tk.BOTH)
        
        # 图像显示区域
        self.image_label = Label(left_frame, text="原始图像")
        self.image_label.pack()
        
        self.canvas_original = Label(left_frame)
        self.canvas_original.pack(padx=10, pady=5)
        
        # 为原始图像添加点击事件
        self.canvas_original.bind("<Button-1>", self.pick_color)
        
        self.mask_label = Label(left_frame, text="蒙版预览")
        self.mask_label.pack()
        
        self.canvas_mask = Label(left_frame)
        self.canvas_mask.pack(padx=10, pady=5)
        
        # 控制面板
        control_frame = Frame(right_frame)
        control_frame.pack(fill=tk.BOTH, padx=10, pady=10)
        
        # 文件操作按钮
        file_frame = Frame(control_frame)
        file_frame.pack(fill=tk.X, pady=5)
        
        self.load_btn = Button(file_frame, text="加载图像", command=self.load_image)
        self.load_btn.pack(side=tk.LEFT, padx=5)
        
        self.save_btn = Button(file_frame, text="保存蒙版", command=self.save_mask)
        self.save_btn.pack(side=tk.LEFT, padx=5)
        
        self.process_btn = Button(file_frame, text="生成蒙版", command=self.update_preview)
        self.process_btn.pack(side=tk.LEFT, padx=5)
        
        # 增加颜色拾取开关按钮
        self.color_pick_btn = Button(file_frame, text="颜色拾取", command=self.toggle_color_pick)
        self.color_pick_btn.pack(side=tk.LEFT, padx=5)
        
        # 参数调整区域
        params_frame = Frame(control_frame)
        params_frame.pack(fill=tk.BOTH, expand=True, pady=10)
        
        # 添加颜色选择器部分
        color_section_frame = Frame(params_frame)
        color_section_frame.pack(fill=tk.X, pady=5)
        Label(color_section_frame, text="颜色拾取", font=("Arial", 10, "bold")).pack(anchor="w")
        
        # 显示选中颜色的标签
        color_display_frame = Frame(params_frame)
        color_display_frame.pack(fill=tk.X, pady=2)
        Label(color_display_frame, text="选中颜色:").pack(side=tk.LEFT)
        self.color_canvas = Canvas(color_display_frame, width=30, height=20, bg="#000000")
        self.color_canvas.pack(side=tk.LEFT, padx=5)
        self.color_text = StringVar(value="RGB: (0,0,0)")
        Label(color_display_frame, textvariable=self.color_text).pack(side=tk.LEFT)
        
        # 颜色容差滑块
        tolerance_frame = Frame(params_frame)
        tolerance_frame.pack(fill=tk.X, pady=2)
        Label(tolerance_frame, text="颜色容差:").pack(side=tk.LEFT)
        self.tolerance_var = IntVar(value=30)
        Scale(tolerance_frame, from_=0, to=100, orient=tk.HORIZONTAL,
              variable=self.tolerance_var, command=self.update_color_mask).pack(side=tk.LEFT, fill=tk.X, expand=True)
        
        # 分隔线
        separator = Frame(params_frame, height=2, bg="gray")
        separator.pack(fill=tk.X, pady=10)
        
        # 预处理参数
        Label(params_frame, text="图像预处理参数").pack(anchor='w')
        
        # 模糊大小
        blur_frame = Frame(params_frame)
        blur_frame.pack(fill=tk.X, pady=2)
        Label(blur_frame, text="模糊大小:").pack(side=tk.LEFT)
        self.blur_var = IntVar(value=5)
        Scale(blur_frame, from_=0, to=21, orient=tk.HORIZONTAL, 
              variable=self.blur_var, resolution=2).pack(side=tk.LEFT, fill=tk.X, expand=True)
        
        # 亮度
        bright_frame = Frame(params_frame)
        bright_frame.pack(fill=tk.X, pady=2)
        Label(bright_frame, text="亮度调整:").pack(side=tk.LEFT)
        self.bright_var = IntVar(value=0)
        Scale(bright_frame, from_=-50, to=50, orient=tk.HORIZONTAL, 
              variable=self.bright_var).pack(side=tk.LEFT, fill=tk.X, expand=True)
        
        # 对比度
        contrast_frame = Frame(params_frame)
        contrast_frame.pack(fill=tk.X, pady=2)
        Label(contrast_frame, text="对比度:").pack(side=tk.LEFT)
        self.contrast_var = IntVar(value=10)
        Scale(contrast_frame, from_=1, to=30, orient=tk.HORIZONTAL, 
              variable=self.contrast_var).pack(side=tk.LEFT, fill=tk.X, expand=True)
        
        # CLAHE参数
        clahe_frame = Frame(params_frame)
        clahe_frame.pack(fill=tk.X, pady=2)
        self.clahe_var = IntVar(value=1)
        tk.Checkbutton(clahe_frame, text="使用CLAHE增强", 
                     variable=self.clahe_var).pack(side=tk.LEFT)
        
        clip_frame = Frame(params_frame)
        clip_frame.pack(fill=tk.X, pady=2)
        Label(clip_frame, text="CLAHE强度:").pack(side=tk.LEFT)
        self.clip_var = IntVar(value=20)
        Scale(clip_frame, from_=1, to=50, orient=tk.HORIZONTAL, 
              variable=self.clip_var).pack(side=tk.LEFT, fill=tk.X, expand=True)
        
        # 阈值处理参数
        Label(params_frame, text="\n阈值处理参数").pack(anchor='w')
        
        thresh_frame = Frame(params_frame)
        thresh_frame.pack(fill=tk.X, pady=2)
        Label(thresh_frame, text="阈值方法:").pack(side=tk.LEFT)
        self.thresh_method = StringVar(value="adaptive")
        tk.Radiobutton(thresh_frame, text="自适应阈值", 
                      variable=self.thresh_method, value="adaptive").pack(side=tk.LEFT)
        tk.Radiobutton(thresh_frame, text="Otsu阈值", 
                      variable=self.thresh_method, value="otsu").pack(side=tk.LEFT)
        
        block_frame = Frame(params_frame)
        block_frame.pack(fill=tk.X, pady=2)
        Label(block_frame, text="Block大小:").pack(side=tk.LEFT)
        self.block_var = IntVar(value=11)
        Scale(block_frame, from_=3, to=51, orient=tk.HORIZONTAL, 
              variable=self.block_var, resolution=2).pack(side=tk.LEFT, fill=tk.X, expand=True)
        
        c_frame = Frame(params_frame)
        c_frame.pack(fill=tk.X, pady=2)
        Label(c_frame, text="常数C:").pack(side=tk.LEFT)
        self.c_var = IntVar(value=2)
        Scale(c_frame, from_=-10, to=10, orient=tk.HORIZONTAL, 
              variable=self.c_var).pack(side=tk.LEFT, fill=tk.X, expand=True)
        
        # 形态学操作参数
        Label(params_frame, text="\n形态学处理参数").pack(anchor='w')
        
        morph_frame = Frame(params_frame)
        morph_frame.pack(fill=tk.X, pady=2)
        Label(morph_frame, text="操作类型:").pack(side=tk.LEFT)
        self.morph_op = StringVar(value="close")
        ops = [("膨胀", "dilate"), ("腐蚀", "erode"), 
               ("开运算", "open"), ("闭运算", "close")]
        for text, value in ops:
            tk.Radiobutton(morph_frame, text=text, 
                          variable=self.morph_op, value=value).pack(side=tk.LEFT)
        
        kernel_frame = Frame(params_frame)
        kernel_frame.pack(fill=tk.X, pady=2)
        Label(kernel_frame, text="核大小:").pack(side=tk.LEFT)
        self.kernel_var = IntVar(value=3)
        Scale(kernel_frame, from_=1, to=21, orient=tk.HORIZONTAL, 
              variable=self.kernel_var, resolution=2).pack(side=tk.LEFT, fill=tk.X, expand=True)
        
        # 连通区域过滤
        Label(params_frame, text="\n连通区域过滤").pack(anchor='w')
        
        area_frame = Frame(params_frame)
        area_frame.pack(fill=tk.X, pady=2)
        Label(area_frame, text="最小面积:").pack(side=tk.LEFT)
        self.area_var = IntVar(value=100)
        Scale(area_frame, from_=0, to=1000, orient=tk.HORIZONTAL, 
              variable=self.area_var).pack(side=tk.LEFT, fill=tk.X, expand=True)
        
        # 状态信息
        self.status_var = StringVar(value="请加载图像")
        status_bar = Label(self.root, textvariable=self.status_var, bd=1, relief=tk.SUNKEN, anchor=tk.W)
        status_bar.pack(side=tk.BOTTOM, fill=tk.X)
    
    def toggle_color_pick(self):
        """切换颜色拾取模式"""
        self.color_pick_mode = not self.color_pick_mode
        if self.color_pick_mode:
            self.color_pick_btn.config(relief=tk.SUNKEN, bg="#aaaaaa")
            self.status_var.set("颜色拾取模式: 点击图像选择颜色")
        else:
            self.color_pick_btn.config(relief=tk.RAISED, bg=self.root.cget("bg"))
            self.status_var.set("颜色拾取模式已关闭")
    
    def pick_color(self, event):
        """从图像中拾取颜色"""
        if not self.color_pick_mode or not hasattr(self.processor, 'original') or self.processor.original is None:
            return
            
        # 获取点击位置
        x, y = event.x, event.y
        
        # 检查图像尺寸，计算实际位置
        display_img = self.canvas_original.image  # 这是PIL图像对象
        if display_img:
            # 获取显示图像的尺寸
            display_width = display_img.width()
            display_height = display_img.height()
            
            # 获取原始图像尺寸
            original_height, original_width = self.processor.original.shape[:2]
            
            # 计算缩放比例
            scale_x = original_width / display_width
            scale_y = original_height / display_height
            
            # 计算在原图中的位置
            original_x = min(max(int(x * scale_x), 0), original_width - 1)
            original_y = min(max(int(y * scale_y), 0), original_height - 1)
            
            # 获取颜色（BGR格式）
            color = self.processor.original[original_y, original_x]
            b, g, r = color
            
            # 更新选中的颜色
            self.processor.selected_color = (int(b), int(g), int(r))
            
            # 更新颜色显示
            hex_color = f"#{r:02x}{g:02x}{b:02x}"
            self.color_canvas.config(bg=hex_color)
            self.color_text.set(f"RGB: ({r},{g},{b})")
            
            # 创建基于颜色的蒙版
            self.update_color_mask()
            
            self.status_var.set(f"已选择颜色: RGB({r},{g},{b})")
    
    def update_color_mask(self, *args):
        """基于选中的颜色更新蒙版预览"""
        if not hasattr(self.processor, 'selected_color') or self.processor.selected_color is None:
            return
            
        try:
            self.status_var.set("处理中...")
            self.root.update()
            
            # 获取容差值
            tolerance = self.tolerance_var.get()
            
            # 生成基于颜色的蒙版
            start_time = time.time()
            mask = self.processor.create_color_mask(self.processor.selected_color, tolerance)
            process_time = time.time() - start_time
            
            # 显示蒙版
            self.display_image(mask, self.canvas_mask)
            
            self.status_var.set(f"颜色蒙版已生成 (处理时间: {process_time:.2f}秒)")
        except Exception as e:
            self.status_var.set(f"颜色处理错误: {str(e)}")
    
    def load_image(self):
        """加载图像文件"""
        file_path = filedialog.askopenfilename(
            filetypes=[("图像文件", "*.jpg *.jpeg *.png *.bmp")],
            initialdir=os.path.dirname(os.path.dirname(os.getcwd()))
        )
        
        if file_path:
            try:
                self.image_path = file_path
                self.processor.load_image(file_path)
                self.display_image(self.processor.image, self.canvas_original)
                self.status_var.set(f"已加载图像: {os.path.basename(file_path)}")
                self.update_preview()
            except Exception as e:
                self.status_var.set(f"错误: {str(e)}")
    
    def update_preview(self):
        """更新蒙版预览"""
        if not hasattr(self.processor, 'original') or self.processor.original is None:
            self.status_var.set("请先加载图像")
            return
            
        try:
            self.status_var.set("处理中...")
            self.root.update()
            
            # 获取当前所有参数
            params = {
                'blur_size': self.blur_var.get(),
                'brightness': self.bright_var.get(),
                'contrast': self.contrast_var.get() / 10.0,
                'use_clahe': bool(self.clahe_var.get()),
                'clip_limit': self.clip_var.get() / 10.0,
                'thresh_method': self.thresh_method.get(),
                'block_size': self.block_var.get(),
                'c': self.c_var.get(),
                'morph_op': self.morph_op.get(),
                'kernel_size': self.kernel_var.get(),
                'min_area': self.area_var.get()
            }
            
            # 生成蒙版
            start_time = time.time()
            mask = self.processor.create_mask(**params)
            process_time = time.time() - start_time
            
            # 显示蒙版
            self.display_image(mask, self.canvas_mask)
            
            self.status_var.set(f"蒙版已生成 (处理时间: {process_time:.2f}秒)")
        except Exception as e:
            self.status_var.set(f"处理错误: {str(e)}")
    
    def save_mask(self):
        """保存蒙版图像"""
        if not hasattr(self.processor, 'mask') or self.processor.mask is None:
            self.status_var.set("请先生成蒙版")
            return
            
        # 构建默认文件名
        if self.image_path:
            directory = os.path.dirname(self.image_path)
            basename = os.path.splitext(os.path.basename(self.image_path))[0]
            default_filename = f"{basename}_mask.jpg"
            default_path = os.path.join(directory, default_filename)
        else:
            default_path = "mask.jpg"
            
        file_path = filedialog.asksaveasfilename(
            defaultextension=".jpg",
            filetypes=[("JPEG 文件", "*.jpg"), ("所有文件", "*.*")],
            initialfile=default_path
        )
        
        if file_path:
            try:
                self.processor.save_mask(file_path)
                self.status_var.set(f"蒙版已保存至: {os.path.basename(file_path)}")
            except Exception as e:
                self.status_var.set(f"保存错误: {str(e)}")
    
    def display_image(self, cv_image, canvas):
        """在画布上显示OpenCV图像"""
        # 确保图像是灰度或BGR格式
        if len(cv_image.shape) == 2:
            # 灰度图像转RGB
            display_image = cv2.cvtColor(cv_image, cv2.COLOR_GRAY2RGB)
        else:
            # BGR转RGB
            display_image = cv2.cvtColor(cv_image, cv2.COLOR_BGR2RGB)
            
        # 调整图像大小以适应界面
        max_height = 300
        height, width = display_image.shape[:2]
        
        if height > max_height:
            scale = max_height / height
            new_width = int(width * scale)
            display_image = cv2.resize(display_image, (new_width, max_height))
            
        # 转换为Tkinter可显示的格式
        pil_image = Image.fromarray(display_image)
        tk_image = ImageTk.PhotoImage(image=pil_image)
        
        # 更新画布
        canvas.config(image=tk_image)
        canvas.image = tk_image  # 保持引用以防止被垃圾回收


def main():
    """主函数"""
    root = tk.Tk()
    app = MaskGeneratorUI(root)
    root.mainloop()


if __name__ == "__main__":
    main()
